import audioop
from enum import Enum
from typing import Type
from urllib import parse
import pyaudio
import wave
import os
import requests
import urllib.request
import base64
import json
import rclpy
from rclpy.parameter import Parameter
from rclpy.node import Node
from std_msgs.msg import String
from std_msgs.msg import Bool


class Record(Enum):
    LEVEL_DETECTION = 0
    TIME_QUANTUM = 1


# 路径
PACKAGE_NAME = 'robot_voice'
ASSETS_DIR = os.path.join(os.path.dirname(
    os.path.abspath(__file__)), "../../../../share", PACKAGE_NAME)
# pyaudio参数
CHUNK = 1024
FORMAT = pyaudio.paInt16
CHANNELS = 1                # 声道数
RATE = 16000                # 采样率
RECORD_SECONDS = 5
WAVE_OUTPUT_FILENAME = os.path.join(ASSETS_DIR, "output.wav")
# 阈值参数
VOICE_THRESHOLD = 1000
VOICE_MONITORING = 50
RECORD_TYPE = Record.LEVEL_DETECTION


class SpeechRecognition(Node):
    def __init__(self, name):
        super().__init__(name)
        self.get_logger().info("启动节点: %s" % name)

        # 服务器地址
        self.declare_parameter('host', '127.0.0.1')
        self.declare_parameter('port', 8090)
        
        self.get_logger().info("服务器地址: %s" % f'http://{self.get_parameter("host").value}:{self.get_parameter("port").value}')

        self.tts_publisher = self.create_publisher(
            String, "tts", 10)

        self.wake_up_subscribe = self.create_subscription(
            Bool, "wake_up", self.wake_up_callback, 1)

        self.declare_parameters

    def wake_up_callback(self, message):
        if message.data:
            self.get_logger().info("测试")
            self.get_audio()
            self.recognite()

    def get_audio(self):
        pa = pyaudio.PyAudio()
        stream = pa.open(format=FORMAT,
                         channels=CHANNELS,
                         rate=RATE,
                         input=True,
                         frames_per_buffer=CHUNK)
        self.get_logger().info("开始录音...")
        frames = []

        if RECORD_TYPE == Record.LEVEL_DETECTION:
            # 检测声音强弱
            low_audio_flag = 0
            while True:
                stream_data = stream.read(CHUNK)
                rms = audioop.rms(stream_data, 2)

                self.get_logger().debug("音量电平: " + str(rms))

                low_audio_flag = 0 if rms > VOICE_THRESHOLD else low_audio_flag + 1
                # 连续 VOICE_MONITORING 次采样都是小音量，则认为没有音频
                if low_audio_flag > VOICE_MONITORING:
                    break
                frames.append(stream_data)

        elif RECORD_TYPE == Record.TIME_QUANTUM:
            # 录制指定时间
            for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
                stream_data = stream.read(CHUNK)
                frames.append(stream_data)

        self.get_logger().info("录音结束...")

        stream.stop_stream()
        stream.close()
        pa.terminate()

        wf = wave.open(WAVE_OUTPUT_FILENAME, 'wb')
        wf.setnchannels(CHANNELS)
        wf.setsampwidth(pa.get_sample_size(FORMAT))
        wf.setframerate(RATE)
        wf.writeframes(b''.join(frames))
        wf.close()

    def paddlespeech_request(self, url, data):
        """ 
        构造PaddleSpeech所需的post请求
        """
        res = requests.post(
            url=url,
            data=json.dumps(data)
        )

        if res.status_code == 200:
            res = res.json()
        else:
            self.get_logger().info("请求失败，错误代码: " + str(res.status_code))
            res = None
        return res

    def recognite(self):
        self.get_logger().info("开始识别...")

        # 将wav转成base64
        with open(WAVE_OUTPUT_FILENAME, 'rb') as f:
            base64_bytes = base64.b64encode(f.read())
            base64_string = base64_bytes.decode('utf-8')

        data = {
            "audio": base64_string,
            "audio_format": "wav",
            "sample_rate": 16000,
            "lang": "zh_cn",
            "punc": 0
        }

        url = f'http://{self.get_parameter("host").value}:{self.get_parameter("port").value}/paddlespeech/asr'
        res = self.paddlespeech_request(url, data)
        success = res["success"]
        text = res["result"]["transcription"]

        if success:
            self.get_logger().info("识别结果: " + text)
            self.recognite_success(text)
            self.punctuation(text)
        else:
            self.get_logger().info("识别失败!")

    def punctuation(self, text):
        data = {
            "text": text
        }

        url = f'http://{self.get_parameter("host").value}:{self.get_parameter("port").value}/paddlespeech/text'
        res = self.paddlespeech_request(url, data)
        success = res["success"]
        text = res["result"]["punc_text"]
        if success:
            self.get_logger().info("标点解析结果: " + text)
            self.punctuation_success(text)
        else:
            self.get_logger().info("标点解析失败!")

    def recognite_success(self, text):
        return

    def punctuation_success(self, text):
        match text:
            case _:
                reply = self.qingyunke(text)
                if reply != "":
                    self.tts_publisher.publish(String(data=reply))

    def qingyunke(self, message):
        url = 'http://api.qingyunke.com/api.php?key=free&appid=0&msg={}'.format(
            urllib.parse.quote(str(message)))
        html = requests.get(url)
        res = html.json()
        if res["result"] == 0:
            self.get_logger().info("青云客回复获取成功: " + res["content"])
            return res["content"]
        else:
            self.get_logger().info("青云客回复获取失败...")
            return ""


def main(args=None):
    rclpy.init(args=args)
    node = SpeechRecognition("speech_recognition")
    rclpy.spin(node)
    rclpy.shutdown()
